"""
图片生成工具
用于预生成不同大小的测试图片，避免测试执行时的动态生成开销
优化版本：支持PNG、JPEG三种格式，快速生成，低CPU消耗
"""
import os
import sys
from pathlib import Path
from PIL import Image
from typing import List, Dict, Tuple
import random


class ImageGenerator:
    """图片生成器类 - 优化版本，支持三种常见格式"""
    
    # 支持的图片格式及其默认参数
    SUPPORTED_FORMATS = {
        'PNG': {
            'extension': 'png',
            'save_params': {'format': 'PNG', 'optimize': False},
            'compression_ratio': 50,  # 纯色PNG压缩率很高
            'quality_param': None
        },
        'JPEG': {
            'extension': 'jpeg', 
            'save_params': {'format': 'JPEG'},
            'compression_ratio': 20,  # JPEG对纯色压缩较好
            'quality_param': 'quality'
        },
        'JPG': {
            'extension': 'jpg',
            'save_params': {'format': 'JPEG'},
            'compression_ratio': 20,  # JPG实际上就是JPEG
            'quality_param': 'quality'
        }
    }
    
    def __init__(self, output_dir: str = None):
        """
        初始化图片生成器
        
        Args:
            output_dir: 输出目录，默认为 cases/picture-file
        """
        if output_dir is None:
            # 获取项目根目录
            current_dir = Path(__file__).parent
            project_root = current_dir.parent
            output_dir = project_root / "cases" / "picture-file"
        
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
    def _validate_format(self, format_type: str) -> str:
        """
        验证并标准化图片格式
        
        Args:
            format_type: 图片格式字符串
            
        Returns:
            标准化的格式名称
            
        Raises:
            ValueError: 不支持的格式
        """
        format_upper = format_type.upper()
        if format_upper not in self.SUPPORTED_FORMATS:
            supported = list(self.SUPPORTED_FORMATS.keys())
            raise ValueError(f"不支持的格式 '{format_type}'。支持的格式: {supported}")
        return format_upper
    
    def _create_simple_image(self, width: int, height: int) -> Image.Image:
        """
        快速创建简单图像 - 使用纯色块，极速生成
        
        Args:
            width: 图片宽度
            height: 图片高度
            
        Returns:
            PIL Image对象
        """
        # 使用随机颜色创建纯色图像，这是最快的方式
        color = (
            random.randint(50, 200),  # R
            random.randint(50, 200),  # G  
            random.randint(50, 200)   # B
        )
        return Image.new('RGB', (width, height), color)
    
    def _estimate_initial_size(self, target_size_mb: float, format_type: str) -> int:
        """
        智能估算初始图片尺寸，针对不同格式优化
        
        Args:
            target_size_mb: 目标文件大小（MB）
            format_type: 图片格式
            
        Returns:
            估算的图片边长
        """
        target_size_bytes = target_size_mb * 1024 * 1024
        format_info = self.SUPPORTED_FORMATS[format_type]
        
        # 根据格式的压缩特性估算像素数
        estimated_pixels = target_size_bytes * format_info['compression_ratio']
        side_length = int(estimated_pixels ** 0.5)
        
        return max(100, side_length)  # 最小100像素
    
    def _get_save_params(self, format_type: str, quality: int = 85) -> Dict:
        """
        获取保存参数
        
        Args:
            format_type: 图片格式
            quality: 质量参数 (1-100)
            
        Returns:
            保存参数字典
        """
        format_info = self.SUPPORTED_FORMATS[format_type]
        save_params = format_info['save_params'].copy()
        
        # 添加质量参数（如果格式支持）
        if format_info['quality_param'] and quality is not None:
            save_params[format_info['quality_param']] = quality
            
        return save_params
    
    def _cleanup_failed_file(self, filepath: Path):
        """
        清理生成失败的文件
        
        Args:
            filepath: 要清理的文件路径
        """
        try:
            if filepath.exists():
                filepath.unlink()
                print(f"  已清理失败的文件: {filepath.name}")
        except Exception as e:
            print(f"  清理文件失败: {e}")
    
    def generate_image_by_size(self, target_size_mb: float, format_type: str = 'PNG', 
                              quality: int = 85) -> str:
        """
        快速生成指定大小的图片 - 支持三种格式的优化版本
        
        Args:
            target_size_mb: 目标文件大小（MB）
            format_type: 图片格式 ('PNG', 'JPEG', 'JPG')
            quality: 质量参数 (1-100)，PNG忽略此参数
            
        Returns:
            生成的图片文件路径
            
        Raises:
            ValueError: 不支持的格式
        """
        # 验证格式
        format_type = self._validate_format(format_type)
        target_size_bytes = int(target_size_mb * 1024 * 1024)
        
        # 获取格式信息
        format_info = self.SUPPORTED_FORMATS[format_type]
        extension = format_info['extension']
        
        # 快速估算初始尺寸
        side_length = self._estimate_initial_size(target_size_mb, format_type)
        
        # 文件名 - 修改为test_images_xxMB格式
        filename = f"test_images_{target_size_mb:.1f}MB.{extension}"
        filepath = self.output_dir / filename
        
        print(f"正在生成 {target_size_mb}MB 的 {format_type} 图片...")
        
        # 快速迭代调整 - 减少迭代次数
        max_iterations = 10
        tolerance = 0.15  # 15%误差容忍度
        
        try:
            for iteration in range(max_iterations):
                # 快速生成纯色图像
                image = self._create_simple_image(side_length, side_length)
                
                # 获取保存参数并保存
                save_params = self._get_save_params(format_type, quality)
                try:
                    image.save(filepath, **save_params)
                except Exception as e:
                    self._cleanup_failed_file(filepath)
                    raise RuntimeError(f"保存 {format_type} 格式图片失败: {e}")
                
                current_size = os.path.getsize(filepath)
                size_ratio = current_size / target_size_bytes
                
                print(f"  迭代 {iteration + 1}: {side_length}x{side_length} -> "
                      f"{current_size / (1024*1024):.2f}MB (目标: {target_size_mb}MB)")
                
                # 检查是否达到目标
                if abs(size_ratio - 1.0) <= tolerance:
                    print(f"✓ 生成完成: {filepath}")
                    print(f"  文件大小: {current_size / (1024*1024):.2f}MB")
                    return str(filepath)
                
                # 智能调整尺寸
                if size_ratio < 0.5:  # 文件太小，大幅增加
                    side_length = int(side_length * 1.5)
                elif size_ratio < 1.0:  # 文件偏小，适度增加
                    side_length = int(side_length * 1.2)
                elif size_ratio > 2.0:  # 文件太大，大幅减少
                    side_length = int(side_length * 0.6)
                else:  # 文件偏大，适度减少
                    side_length = int(side_length * 0.85)
                
                # 确保最小尺寸
                side_length = max(50, side_length)
            
            # 返回最接近的结果
            current_size = os.path.getsize(filepath)
            print(f"✓ 生成完成: {filepath}")
            print(f"  文件大小: {current_size / (1024*1024):.2f}MB (目标: {target_size_mb}MB)")
            
            return str(filepath)
            
        except Exception as e:
            # 生成失败时清理文件
            self._cleanup_failed_file(filepath)
            raise e
    
    def generate_standard_test_images(self, include_formats: List[str] = None) -> Dict[str, str]:
        """
        快速生成标准测试图片集合
        
        Args:
            include_formats: 要包含的格式列表，默认为所有支持的格式
            
        Returns:
            生成的图片文件路径字典
        """
        if include_formats is None:
            include_formats = list(self.SUPPORTED_FORMATS.keys())
        else:
            # 验证格式
            include_formats = [self._validate_format(fmt) for fmt in include_formats]
        
        # 修改为0.5, 1, 5, 10, 20MB
        sizes = [0.5, 1, 5, 10, 20]  # MB
        qualities = {'JPEG': [85], 'JPG': [85], 'PNG': [None]}  # 简化质量配置
        
        generated_files = {}
        failed_files = []
        
        print("开始快速生成标准测试图片...")
        print("=" * 60)
        
        for size_mb in sizes:
            for format_type in include_formats:
                format_qualities = qualities.get(format_type, [85])
                
                for quality in format_qualities:
                    try:
                        filepath = self.generate_image_by_size(size_mb, format_type, quality)
                        
                        # 生成唯一的key
                        key = f"{size_mb}MB_{format_type}"
                        generated_files[key] = filepath
                        print()  # 空行分隔
                        
                    except Exception as e:
                        quality_info = f" (质量: {quality})" if quality is not None else ""
                        error_msg = f"生成 {size_mb}MB {format_type}{quality_info} 图片失败: {e}"
                        print(f"❌ {error_msg}")
                        failed_files.append(error_msg)
        
        print("=" * 60)
        print(f"✓ 完成! 共生成 {len(generated_files)} 个测试图片")
        if failed_files:
            print(f"⚠️  失败 {len(failed_files)} 个图片:")
            for failed in failed_files:
                print(f"    {failed}")
        print(f"  输出目录: {self.output_dir}")
        
        # 显示生成的文件统计
        format_counts = {}
        for key in generated_files.keys():
            fmt = key.split('_')[1]  # 提取格式名
            format_counts[fmt] = format_counts.get(fmt, 0) + 1
        
        print("  格式统计:")
        for fmt, count in format_counts.items():
            print(f"    {fmt}: {count} 个文件")
        
        return generated_files
    
    def get_supported_formats(self) -> List[str]:
        """
        获取当前支持的图片格式列表
        
        Returns:
            支持的格式列表
        """
        return list(self.SUPPORTED_FORMATS.keys())


def main():
    """主函数 - 命令行入口"""
    print("图片生成工具 - 支持PNG、JPEG、JPG格式")
    print("=" * 60)
    
    # 创建生成器
    generator = ImageGenerator()
    
    # 显示支持的格式
    supported_formats = generator.get_supported_formats()
    print(f"支持的格式: {', '.join(supported_formats)}")
    print()
    
    # 生成标准测试图片
    try:
        generated_files = generator.generate_standard_test_images()
        print(f"\n✓ 所有图片已生成到: {generator.output_dir}")
        
        # 显示生成的文件列表
        print("\n生成的文件:")
        for key, filepath in generated_files.items():
            file_size = os.path.getsize(filepath) / (1024 * 1024)
            print(f"  {key}: {file_size:.2f}MB - {Path(filepath).name}")
        
    except Exception as e:
        print(f"❌ 生成过程中出现错误: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()